#!/usr/bin/env python
# -*- coding: utf-8 -*-

#
# This program is free software; you can redistribute it and/or modify  
# it under the terms of the GNU General Public License as published by  
# the Free Software Foundation; either version 2 of the License, or     
# (at your option) any later version.                                   
#                                                                         
# A copy of the license can be found in the license.txt file supplied   
# with this software or at: http://www.gnu.org/copyleft/gpl.html       
#

'''
FORMAT of SSTable :
+------------------------------+
|  KEY \t VALUE \n             |
|  KEY \t VALUE \n             |  DATA
|  ... ... ...                 | SECTION
|  KEY \t VALUE \n             |
+------------------------------+
|  KEY \t OFFSET \t LENGTH \n  |
|  KEY \t OFFSET \t LENGTH \n  |  INDEX
|  ... ... ...                 | SECTION
|  KEY \t OFFSET \t LENGTH \n  |
+------------------------------+
|  KEY \t INDEX-OFFSET \n      |  SCAN
|  KEY \t INDEX-OFFSET \n      |  INDEX
|  ... ... ...                 | SECTION
|  KEY \t INDEX-OFFSET \n      |
+------------------------------+
|  BLOOM-FILTER \n             | FILTER-SECTION
+------------------------------+
|  PTR-INDEX-SECTION           |   END
|  PTR-SCAN-INDEX-SECTION      |  INDEX
|  PTR-FILTER-SECTION          | SECTION
|  END_OF_SSTABLE              |
+------------------------------+
'''

import cStringIO

from bisect import bisect_left
from jutebag import BloomFilter


SEEK_SET = 0
SEEK_CUR = 1
SEEK_END = 2

class   SSTable(object) :
    '''
    SSTable is a WRITE-ONCE table. If it is opened for writing, 
    content of SSTable file would be truncated. 
    When a SSTable file closes, it should serves READ-ONLY.
    '''

    KV_SEPARATOR = "\t"
    KV_END_TAG = "\n"
    KV_SEPARATOR_LENGTH = 1
    KV_END_TAG_LENGTH = 1

    INDEX_INTERVAL = 128

    END_TAG = 'END_OF_SSTABLE'
    END_INDEX_LENGTH = 16 * 4

    def __init__(self) :
        pass
        ## __init__()

    ## class SSTable

class   SSTableReader(SSTable) :
    def __init__(self, filename) :
        self._filename = filename
        self._fp = None
        self._scan_index_list = []
        self._bf = None     # bloom filter

        self.__open()
        ## __init__()

    def __del__(self) :
        self.close()
        ## __del__()

    def __open(self) :
        self._fp = file(self._filename, 'r')

        # read END-INDEX
        self.__loadEndIndex()

        # read SCAN-INDEX
        self.__loadScanIndex()

        # read Bloom-Filter
        self.__loadBloomFilter()

        ## __open()

    def __loadEndIndex(self) :
        self._fp.seek(-SSTable.END_INDEX_LENGTH, SEEK_END)
        end_index = self._fp.read(SSTable.END_INDEX_LENGTH)
        end_indexes = end_index.split()

        if len(end_indexes) != 4 :
            raise RuntimeError, 'Invalid END-INDEX "%s"' % end_index
        end_tag = end_indexes[3].strip()
        if end_tag != SSTable.END_TAG :
            raise RuntimeError, 'Invalid END-INDEX "%s"' % end_index

        self._index_section_offset = int(end_indexes[0])
        self._scan_index_section_offset = int(end_indexes[1])
        self._filter_section_offset = int(end_indexes[2])

        #print 'IndexSectionOffset    :',self._index_section_offset
        #print 'ScanIndexSectionOffset:',self._scan_index_section_offset
        #print 'FilterSectionOffset   :',self._filter_section_offset
        ## __loadEndIndex()

    def __loadScanIndex(self) :
        self._scan_index_list = []
        self._fp.seek(self._scan_index_section_offset, SEEK_SET)
        end_of_scan_index_section = self._filter_section_offset
        current_offset = self._fp.tell()
        while current_offset < end_of_scan_index_section :
            line = self._fp.readline()
            key, index_offset = line.strip().split(SSTable.KV_SEPARATOR)
            self._scan_index_list.append( (key, int(index_offset)) )
            current_offset = self._fp.tell()

        self._scan_index_length = len(self._scan_index_list)
        #print self._scan_index_list
        ## __loadScanIndex()

    def __loadBloomFilter(self) :
        self._fp.seek(0, SEEK_END)
        end_of_filter_section = self._fp.tell() - SSTable.END_INDEX_LENGTH
        size_of_filter_section = end_of_filter_section - self._filter_section_offset
        self._fp.seek(self._filter_section_offset, SEEK_SET)
        filter_serialized = self._fp.read( size_of_filter_section )

        bf = BloomFilter()
        self._bf = bf.deserialize( filter_serialized.strip() )
        #print str(self._bf)
        ## __loadBloomFilter()

    def close(self) :
        if self._fp :
            self._fp.close()

        self._scan_index_list = []
        self._bf = None     # bloom filter
        ## close()

    def getApproximateKeyCount(self) :
        count_interval = len(self._scan_index_list)

        return count_interval * SSTable.INDEX_INTERVAL
        ## getApproximateKeyCount()

    def __getIndexScanPosition(self, key) :
        if not self._scan_index_list :
            return -1
        index = bisect_left(self._scan_index_list, (key,0))
        if index >= self._scan_index_length :
            index = self._scan_index_length - 1
        elif key != self._scan_index_list[index][0] :
                index -= 1
        if index < 0 :
            index = 0
        scan_index = self._scan_index_list[ index ]
        return scan_index[1]
        ## __getScanIndexPosition()

    def getPosition(self, key) :
        '''get the (position,length) in the data file to find the given key, or (-1,0) if the key is not present

        @TODO to cache posistion of the key

        @return the (position,length) in the data file to find the given key
        '''
        if not self._fp :
            return -1,0
        if not self._bf :
            return -1,0
        if key not in self._bf :
            return -1,0
        index_scan_start = self.__getIndexScanPosition(key)
        if index_scan_start < 0 :
            return -1,0
        index_interval_count_down = SSTable.INDEX_INTERVAL
        end_of_index_section = self._scan_index_section_offset

        self._fp.seek( index_scan_start, SEEK_SET )
        current_offset = self._fp.tell()
        while index_interval_count_down > 0 and current_offset < end_of_index_section :
            line = self._fp.readline()
            if line.startswith(key) :
                _key,_offset,_length = line.strip().split()
                return int(_offset),int(_length)

            index_interval_count_down -= 1
            current_offset = self._fp.tell()
            # end of while

        return -1,0
        ## getPosition()

    def get(self, key) :
        '''get value of the given key, or none if not present
        '''
        offset,length = self.getPosition(key)
        if offset < 0 :
            return None

        self._fp.seek(offset, SEEK_SET)
        data = self._fp.read(length)
        _key,_value = data.split( SSTable.KV_SEPARATOR )
        if _key != key :
            raise RuntimeError, 'Wrong key "%s" ("%s" wanted)' % (_key, key)
        return _value
        ## get()

    ## class SSTableReader


class   SSTableWriter(SSTable) :
    def __init__(self, filename, key_count) :
        self._filename = filename
        self._key_count = key_count
        self._fp = None

        self.__open()
        ## __init__()

    def __del__(self) :
        self.close()
        ## __del__()

    def __open(self) :
        self._fp = file(self._filename, 'w')

        self._index_count = 0
        self._index_interval = 0
        self._last_written_key = None
        self._fp_index = cStringIO.StringIO()
        self._scan_index_list = []
        self._bf = BloomFilter(self._key_count, 15)
        ## __open()

    def close(self) :
        if self._fp :
            # get base offset of Index Section
            index_section_offset = self._fp.tell()

            # write Index Section
            self._fp.write( self._fp_index.getvalue() )

            # get base offset of Scan Index Section
            scan_index_section_offset = self._fp.tell()

            # write Scan Index Section
            for scan_index in self._scan_index_list :
                _key, _offset = scan_index
                self._fp.write("%s%s%d%s" % (_key,SSTable.KV_SEPARATOR,_offset + index_section_offset,SSTable.KV_END_TAG))

            # get base offset of Bloom Filter Section
            bloom_filter_section_offset = self._fp.tell()

            # write Bloom Filter
            self._fp.write(self._bf.serialize())
            self._fp.write(SSTable.KV_END_TAG)

            # write END-INDEX
            self._fp.write("%15d\n" % index_section_offset)
            self._fp.write("%15d\n" % scan_index_section_offset)
            self._fp.write("%15d\n" % bloom_filter_section_offset)
            self._fp.write("%15s\n" % SSTable.END_TAG)

            self._fp.close()
            self._fp = None

        if self._fp_index :
            self._fp_index.close()
        self._scan_index_list = []
        ## close()

    def append(self, key, data) :
        offset = self._beforeAppend(key)

        self._fp.write(key)
        self._fp.write(SSTable.KV_SEPARATOR)
        self._fp.write(str(data))
        self._fp.write(SSTable.KV_END_TAG)

        length = len(key) + SSTable.KV_SEPARATOR_LENGTH + len(str(data))
        self._afterAppend(key, offset, length)
        ## append()

    def _beforeAppend(self, key) :
        '''make sure that key is append in ASCending order
        
        @return current position of file write-pointer
        '''
        # make sure : key is appended in ASCending order
        if self._last_written_key :
            if key <= self._last_written_key :
                raise RuntimeError, 'NOT in ASCending order: %s (last key: %s)' % (key, self._last_written_key)

        return self._fp.tell()
        ## _beforeAppend()

    def _afterAppend(self, key, offset, length) :
        # check if there are too many keys to append
        if self._index_count > self._key_count :
            raise RuntimeError, 'Too many keys (limit:%d)' % (self._key_count,)

        self._last_written_key = key
        if self._index_interval <= 0 :
            # write scan-index item
            index_offset = self._fp_index.tell()
            self._scan_index_list.append( (key, index_offset) )
            self._index_interval = SSTable.INDEX_INTERVAL
        self._index_interval -= 1
        self._index_count += 1

        # write index item
        self._fp_index.write(key)
        self._fp_index.write(SSTable.KV_SEPARATOR)
        self._fp_index.write(str(offset))
        self._fp_index.write(SSTable.KV_SEPARATOR)
        self._fp_index.write(str(length))
        self._fp_index.write(SSTable.KV_END_TAG)

        # add bloom-filter
        self._bf.add(key)

        ## _afterAppend()

    ## class SSTableWriter

